import pandas as pd

df1 = pd.DataFrame(
    [{"r1": 'a', "r2": 1, "symbol": "s1"},
     {"r1": "b", "r2": 2, "symbol": "s2"},
     {"r1": "c", "r2": 2, "symbol": "s3"},
     {"r1": "a", "r2": 3, "symbol": "s4"}
     ]
)

def get_groupby_first():
    buf_df = df1.sort_values("r2", ascending=False)
    grouped = buf_df.groupby("r1")["symbol"].first()
    print(grouped.to_dict())

    grouped2 = buf_df.groupby("r1")["r2"].last()
    print(grouped2.to_dict())


def qmt_deals_group_test():
    deals_df = pd.read_csv("D:\\learn_and_test_data\\deal_sell_20210111.csv", encoding="GBK")
    deals_df["trade_time"] = pd.to_datetime(deals_df["trade_time"], format="%H%M%S")
    deals_df["minuter"] = deals_df["trade_time"].dt.strftime("%H%M")
    grouped = deals_df.groupby(by=["stock_id", "minuter"])
    volume_sum = grouped["volume"].sum()
    for n, rows in deals_df.iterrows():
        key = (rows["stock_id"], rows["minuter"])
        key_vols = volume_sum[key]
        print(key, key_vols)


if __name__ == "__main__":
    #get_groupby_first()
    qmt_deals_group_test()